# ##################################################################################################
#  Disclaimer                                                                                      #
#  This file is a python3 translation of AutoDockTools (v.1.5.7)                                   #
#  Modifications made by Valdes-Tresanco MS (https://github.com/Valdes-Tresanco-MS)                #
#  Tested by Valdes-Tresanco-MS and Valdes-Tresanco ME                                             #
#  There is no guarantee that it works like the original distribution,                             #
#  but feel free to tell us if you get any difference to correct the code.                         #
#                                                                                                  #
#  Please use this cite the original reference.                                                    #
#  If you think my work helps you, just keep this note intact on your program.                     #
#                                                                                                  #
#  Modification date: 10/5/20 18:04                                                                #
#                                                                                                  #
# ##################################################################################################

import numpy


def EulerAnglesToMat(angles):
    """Builds a rotation matrix from Euler angles given in degrees"""

    from math import pi, cos, sin

    t1 = angles[0] * (pi / 180.0)
    t2 = angles[1] * (pi / 180.0)
    t3 = angles[2] * (pi / 180.0)
    # from Lattman, Meth. Enzymology V. 115 p. 63
    rot = numpy.identity(3).astype('d')
    rot[0][0] = -sin(t1) * cos(t2) * sin(t3) + cos(t1) * cos(t3)
    rot[0][1] = cos(t1) * cos(t2) * sin(t3) + sin(t1) * cos(t3)
    rot[0][2] = sin(t2) * sin(t3)
    rot[1][0] = -sin(t1) * cos(t2) * cos(t3) - cos(t1) * sin(t3)
    rot[1][1] = cos(t1) * cos(t2) * cos(t3) - sin(t1) * sin(t3)
    rot[1][2] = sin(t2) * cos(t3)
    rot[2][0] = sin(t1) * sin(t2)
    rot[2][1] = -cos(t1) * sin(t2)
    rot[2][2] = cos(t2)
    return rot


class Crystal:
    """Class to provide functionalities related to crystal structures"""

    def __init__(self, length, angles):
        """Constructor for Crystal"""

        self.ctof = None  # transpose of Xformation to go from cart. to fract.
        self.ftoc = None  # transpose of Xformation to go from fract. to cart.
        assert len(length) == 3 and type(length[0]) == float
        self.length = length
        assert len(angles) == 3 and type(angles[0]) == float
        self.angles = angles
        self.ctof = self._orthog(self.length, self.angles)  # build ctof
        self.ftoc = self._uinv(numpy.transpose(self.ctof))  # build ftoc

    def _orthog(self, length, angles):
        """
           let u be a 3x3 transformation matrix which relates a new cell
           with orthogonal axes of unit dimensions to the original unit
           cell (crystal system).  
            Aug 5, 1991 changed to XPLOR convention was:
            ut[0][0]=as;
            ut[0][1]=0.;
            ut[0][2]=0.;
            ut[1][0]=bs*cosgas;
            ut[1][1]=1./b;
            ut[1][2]= -(1./tan(al))/b;
            ut[2][0]=cs*cosbes;
            ut[2][1]=0.;
            ut[2][2]=1./(c*sin(al));

            June 1, 1996:  Corrected LFT
            The correct orthogonalization matrix for the default
            PDB convention (Cartesion x along A, y in A-B plane,
            and z along c*) is

            1/a   -cos(ga)/(a sin(ga))   as cosbes
             0          1/(b sin(ga))    bs cosals
             0               0              cs

            and the deorthogonalization matrix is

            a   b cos(ga)   c cos(be)
            0   b sin(ga)   -c sin(be) cosals
            0       0         1/cs

            usage:
            xf = x * ut[0][0] + y * ut[0][1] + z * ut[0][2]
            yf = x * ut[1][0] + y * ut[1][1] + z * ut[1][2]
            zf = x * ut[2][0] + y * ut[2][1] + z * ut[2][2]
            """

        from math import pi, cos, sin, sqrt
        degtor = pi / 180.
        al = angles[0] * degtor
        be = angles[1] * degtor
        ga = angles[2] * degtor
        vol = length[0] * length[1] * length[2] * sqrt(1. - cos(al) * cos(al) - cos(be) * cos(be)
                                                       - cos(ga) * cos(ga) + 2. * (cos(al) * cos(be) * cos(ga)))
        lAs = (length[1] * length[2] * sin(al)) / vol
        bs = (length[0] * length[2] * sin(be)) / vol
        cs = (length[0] * length[1] * sin(ga)) / vol
        cosals = (cos(be) * cos(ga) - cos(al)) / (sin(be) * sin(ga))
        cosbes = (cos(ga) * cos(al) - cos(be)) / (sin(ga) * sin(al))
        cosgas = (cos(al) * cos(be) - cos(ga)) / (sin(al) * sin(be))

        ut = numpy.identity(3).astype('f')
        # sabgs1 = sqrt(1.0-cos(al)*cos(al))
        ut[0][0] = 1. / length[0]
        ut[0][1] = -(cos(ga)) / (sin(ga) * length[0])
        # ut[0][2]= -(cos(ga)*sin(be)*cosals + cos(be)*sin(ga))/ */
        #	(sin(be)*sabgs1*sin(ga)*a) */
        ut[0][2] = lAs * cosbes
        ut[1][0] = 0.0
        ut[1][1] = 1. / (sin(ga) * length[1])
        # ut[1][2]=cosals/(sabgs1*sin(ga)*b)
        ut[1][2] = bs * cosals
        ut[2][0] = 0.0
        ut[2][1] = 0.0
        # ut[2][2]=1.0/(sin(be)*sabgs1*c)
        ut[2][2] = cs
        return numpy.transpose(ut)

    def _uinv(self, mat):
        """Inverts a transformation matrix used to go from cartesian to cell
        coordinates"""

        dmat = numpy.identity(3).astype('d')
        imat = numpy.identity(3).astype('d')

        dmat[0][0] = mat[1][1] * mat[2][2] - mat[2][1] * mat[1][2]
        dmat[1][0] = mat[1][2] * mat[2][0] - mat[1][0] * mat[2][2]
        dmat[2][0] = mat[1][0] * mat[2][1] - mat[1][1] * mat[2][0]
        dmat[0][1] = mat[0][2] * mat[2][1] - mat[0][1] * mat[2][2]
        dmat[1][1] = mat[0][0] * mat[2][2] - mat[0][2] * mat[2][0]
        dmat[2][1] = mat[0][1] * mat[2][0] - mat[0][0] * mat[2][1]
        dmat[0][2] = mat[0][1] * mat[1][2] - mat[0][2] * mat[1][1]
        dmat[1][2] = mat[0][2] * mat[1][0] - mat[0][0] * mat[1][2]
        dmat[2][2] = mat[0][0] * mat[1][1] - mat[0][1] * mat[1][0]
        det = (mat[0][0] * mat[1][1] * mat[2][2] + mat[1][0] * mat[2][1] * mat[0][2] + mat[2][0] * mat[0][1] *
               mat[1][2] - mat[2][2] * mat[0][1] * mat[1][0] - mat[0][0] * mat[2][1] * mat[1][2] - mat[2][0] *
               mat[1][1] * mat[0][2])
        for i in (0, 1, 2):
            for j in (0, 1, 2):
                imat[j][i] = dmat[j][i] / det

        return numpy.transpose(imat)

    def toCartesian(self, coords):
        """Transform coords from fractional coordinates to cartesian space"""

        return numpy.dot(coords, self.ftoc)

    def toFractional(self, coords):
        """Transform coords from cartesian coordinates to fractional space"""

        return numpy.dot(coords, self.ctof)

    def rotate(self, rot, coords):
        """Apply rotation to the coordinates. Rotation can be a 3x3 matrix or
        3 Euler angles"""

        if isinstance(rot[0], float):
            rot = EulerAnglesToMat(rot)
        return numpy.dot(coords, numpy.transpose(rot))

    def translate(self, trans, coords):
        """Apply a translation to the coordinates"""

        return coords + trans


if __name__ == '__main__':
    cryst = Crystal((120.3, 120.3, 78.4), (90., 90., 120.))

    from MolKit import Read

    mol = Read('./1ctt_CA.pdb')[0]
    c = numpy.array(mol.chains.residues.atoms.coords)

    fc = cryst.toFractional(c)
    fc1 = cryst.rotate(((0., 1., 0.), (1., 0., 0.), (0., 0., -1.)), fc)
    fc2 = cryst.toCartesian(fc1)
